#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under theglim Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from numpy import dot

# from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors
from camori_harris import PictureHarris
from camori_aux import *
from quaternion import Quat, random_quaternion

import simplejson

import scipy.optimize
#from fmin_sa import Sampling
from simplex_sa import SimplexSO3

import Image

from plot_aux import *

set_printoptions(precision=7)
    
if __name__ == '__main__':
    tt_total_init = time.clock()
    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()
        set_printoptions(precision=3)

    #################################################################
    ## Load image and initialize pic object
    tt_init = time.clock()

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])
    filename = fileroot+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    # Smooth out
    if ("gaussian_smoothing_factor" in job_params.keys() and 
        job_params["gaussian_smoothing_factor"] > 0):
        for c in range(3):
            imr[:,:,c] = scipy.ndimage.gaussian_filter( imr[:,:,c], \
                          double(job_params["gaussian_smoothing_factor"]))

    ## Creates picture object
    if not 'projection_model' in job_params:
        model = 'harris'
        distortion = 0.0
    else:
        model = job_params['projection_model']

    if model != 'harris':
        raise NotImplementedError
    elif model == 'harris':
        ## Intrinsic parameters
        focal_distance = job_params['focal_distance_guess']
        p_point = array(job_params['principal_point_guess'])
        distortion = job_params['distortion_coefficient_guess']
        pic = PictureHarris(imr, p_point[0], p_point[1], focal_distance, distortion)

    tt_image_initialization = time.clock()-tt_init
    ##
    ##################################################################

    ## Edgel extractor parameters
    gmethod = job_params['edge_detection_filter']
    gspc = job_params['grid_spacing']
    glim = job_params['gradient_threshold']
    gfil_t = job_params['edgel_filter_threshold']
    gfil_d = job_params['edgel_filter_distance']
    initial_trials = job_params['initial_trials']

    ##################################################################
    ## Extract the edgels from the image using the grid mask
    tt_init = time.clock()
    pic.extract_edgels(gspc, glim, method=gmethod)
    tt_edgel_extraction = time.clock()-tt_init
    ##
    ##################################################################

    ##################################################################
    ## Filter the edgels to remove the lonely ones. Impose a kind of
    ## curvature restriction in the edgels, but still not a full
    ## countour extraction...  Uses a kd-tree, but still slow.
    if gfil_t > 0:
        pic.edgel_directional_filter(gfil_t, gfil_d)
    ##
    #################################################################

    ##################################################################
    ## Caculate the edgel normals. The normals are used to calculate
    ## initial trials for the orientation estimation. We used to care
    ## about an error function that was directrly based on these
    ## normals, but not anymore. It's just really for the initial
    ## trials.
    ##
    pic.calculate_edgel_normals()
    ##
    ##################################################################

    ##################################################################
    ## Calculate initial estimate
    tt_init = time.clock()
    qopt,bpk_a,bpk_b,bpk_c = pic.estimate_by_random_search(initial_trials)
    tt_initial_estimate = time.clock() - tt_init
    ##
    ##################################################################

    qopt = qopt.canonical()

    tt_total = time.clock() - tt_total_init

    print framenum,
    print qopt.q[0], qopt.q[1], qopt.q[2], qopt.q[3],
    print pic.edgels.shape[0], 
    print tt_edgel_extraction, tt_initial_estimate, tt_total

    if PlotStuff:
        tt = mgrid[0:2*pi:0.001]

        figure(1)
        title('Extracted edgels')
        plot_edgels(gca(),pic)
        pplot(pic.edgels[bpk_a], 'ro')
        pplot(pic.edgels[bpk_b], 'ro')
        pplot(pic.edgels[bpk_c], 'rs')

        figure(2)
        suptitle(u'Estimativa inicial encontrada')
        imshow(.1+pic.frame*.9/255.)
        pic.plot_vdirs(gca(), 31, qopt)
        axis('equal')
        grid()
        pplot(pic.edgels[bpk_a], 'ro')
        pplot(pic.edgels[bpk_b], 'ro')
        pplot(pic.edgels[bpk_c], 'rs')
        axis([0,pic.Iwidth, pic.Iheight,0])
        savefig('robust_demo1.png')
        savefig('robust_demo1.pdf')

        figure(4)
        na = pic.normals[bpk_a]
        nb = pic.normals[bpk_b]
        nc = pic.normals[bpk_c]
        vp = cross(na, nb)
        vp *= (vp**2).sum()**-0.5
        plot_spheres(qopt, pic.normals)
        ii = [[0,1],[2,1],[0,2]]
        for sp in range(3):
            subplot(2,2,1+sp)
            pplot(na[ii[sp]], 'ro')
            pplot(nb[ii[sp]], 'ro')
            pplot(nc[ii[sp]], 'rs')
            pplot(vp[ii[sp]], 'go')
